#!/usr/bin/python3
# -*- coding: utf-8; tab-width: 4; indent-tabs-mode: t -*-

# Copyright (c) 2005-2014 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import os
import re
import sys
import subprocess
import strict_hdds


def usage():
    print('Usage:')
    print('   strict_hdds show <directory>')
    print('   strict_hdds mount <directory> [--layout-name=name] [--disk-list=dev1,dev2...] [--read-only] [--extra-mount-options-for-root-dev=opt1,opt2...] [--extra-mount-options-for-boot-dev=opt1,opt2...]')
    print('   strict_hdds create-and-mount <directory> <layout-name> [--disk-list=dev1,dev2...] [--read-only] [--extra-mount-options-for-root-dev=opt1,opt2...] [--extra-mount-options-for-boot-dev=opt1,opt2...]')
    print('   strict_hdds add-harddisk <directory> <devpath>')
    print('   strict_hdds remove-harddisk <directory> <devpath>')
    print('   strict_hdds create-swap-file <directory> <size>')
    print('   strict_hdds remove-swap-file <directory>')
    print('   strict_hdds create-snapshot <directory> <snapshot-name>')
    print('   strict_hdds remove-snapshot <directory> <snapshot-name>')
    print('   strict_hdds sync-from-snapshot <directory> <snapshot-name> [--sync-home] [--sync-var]')
    print('   strict_hdds check <directory>')
    print('   strict_hdds umount <directory>')
    print('')
    print('Supported layout(s):')
    for layoutName in strict_hdds.get_supported_storage_layout_names():
        print('   %s' % (layoutName))
    print('')


def die(msg):
    print(msg, file=sys.stderr)
    sys.exit(1)


def doShow():
    if len(sys.argv) < 3:
        die("No mount point directory specified!")
    if _getMountDeviceForPath(sys.argv[2]) is None:
        die("Directory \"%s\" is not mounted!")
    layout = strict_hdds.get_storage_layout(sys.argv[2])
    _show(layout)


def doMount():
    if len(sys.argv) < 3:
        die("No mount point directory specified!")
    kwargs = dict()
    for x in sys.argv[3:]:
        if x.startswith("--layout-name="):
            kwargs["layout_name"] = x[len("--layout-name="):]
        if x.startswith("--disk-list="):
            kwargs["disk_list"] = x[len("--disk-list="):].split(",")
        if x == "--read-only":
            kwargs["read_only"] = True
        if x.startswith("--extra-mount-options-for-root-dev="):
            kwargs["extra_mount_options_for_root_dev"] = x[len("--extra-mount-options-for-root-dev="):]
        if x.startswith("--extra-mount-options-for-boot-dev="):
            kwargs["extra_mount_options_for_boot_dev"] = x[len("--extra-mount-options-for-boot-dev="):]
    layout = strict_hdds.mount_storage_layout(sys.argv[2], **kwargs)
    _show(layout)


def doCreateAndMount():
    if len(sys.argv) < 3:
        die("No mount point directory specified!")
    if len(sys.argv) < 4:
        die("No storage layout name specified!")
    kwargs = dict()
    for x in sys.argv[4:]:
        if x.startswith("--disk-list="):
            kwargs["disk_list"] = x[len("--disk-list="):].split(",")
        if x == "--read-only":
            kwargs["read_only"] = True
        if x.startswith("--extra-mount-options-for-root-dev="):
            kwargs["extra_mount_options_for_root_dev"] = x[len("--extra-mount-options-for-root-dev="):]
        if x.startswith("--extra-mount-options-for-boot-dev="):
            kwargs["extra_mount_options_for_boot_dev"] = x[len("--extra-mount-options-for-boot-dev="):]
    layout = strict_hdds.create_and_mount_storage_layout(sys.argv[3], sys.argv[2], **kwargs)
    _show(layout)


def doAddHarddisk():
    if len(sys.argv) < 3:
        die("No mount point directory specified!")
    if len(sys.argv) < 4:
        die("No harddisk device file specified!")
    layout = strict_hdds.get_storage_layout(sys.argv[2])
    if not hasattr(layout, "add_disk"):
        die("Operation is not supported by layout \"%s\"." % (layout.name))
    layout.add_disk(sys.argv[3])


def doRemoveHarddisk():
    if len(sys.argv) < 3:
        die("No mount point directory specified!")
    if len(sys.argv) < 4:
        die("No harddisk device file specified!")
    layout = strict_hdds.get_storage_layout(sys.argv[2])
    if not hasattr(layout, "remove_disk"):
        die("Operation is not supported by storage layout \"%s\"." % (layout.name))
    layout.remove_disk(sys.argv[3])


def doCreateSwapFile():
    if len(sys.argv) < 3:
        die("No mount point directory specified!")

    if len(sys.argv) < 4:
        die('No size (eg: "4G") specified!')

    sizeInGb = None
    if True:
        m = re.fullmatch("([0-9]+)G", sys.argv[3])
        if m is None:
            die('Invalid size (eg: "4G") specified!')
        sizeInGb = int(m.group(1))

    layout = strict_hdds.get_storage_layout(sys.argv[2])
    if not hasattr(layout, "create_swap_file"):
        die("Operation is not supported by layout \"%s\"." % (layout.name))
    layout.create_swap_file(sizeInGb)


def doRemoveSwapFile():
    if len(sys.argv) < 3:
        die("No mount point directory specified!")
    layout = strict_hdds.get_storage_layout(sys.argv[2])
    if not hasattr(layout, "remove_swap_file"):
        die("Operation is not supported by layout \"%s\"." % (layout.name))
    layout.remove_swap_file()


def doCreateSnapshot():
    if len(sys.argv) < 3:
        die("No mount point directory specified!")
    if len(sys.argv) < 4:
        die("No snapshot name specified!")
    layout = strict_hdds.get_storage_layout(sys.argv[2])
    if not hasattr(layout, "create_snapshot"):
        die("Operation is not supported by layout \"%s\"." % (layout.name))
    layout.create_snapshot(sys.argv[3])


def doRemoveSnapshot():
    if len(sys.argv) < 3:
        die("No mount point directory specified!")
    if len(sys.argv) < 4:
        die("No snapshot name specified!")
    layout = strict_hdds.get_storage_layout(sys.argv[2])
    if not hasattr(layout, "remove_snapshot"):
        die("Operation is not supported by layout \"%s\"." % (layout.name))
    layout.remove_snapshot(sys.argv[3])


def doSyncFromSnapshot():
    if len(sys.argv) < 3:
        die("No mount point directory specified!")
    if len(sys.argv) < 4:
        die("No snapshot name specified!")
    kwargs = dict()
    for x in sys.argv[4:]:
        if x == "--sync-home":
            kwargs["home"] = True
        if x == "--sync-var":
            kwargs["var"] = True
    layout = strict_hdds.get_storage_layout(sys.argv[2])
    if not hasattr(layout, "sync_from_snapshot"):
        die("Operation is not supported by layout \"%s\"." % (layout.name))
    layout.sync_from_snapshot(sys.argv[3], **kwargs)


def doCheck():
    if len(sys.argv) < 3:
        die("No mount point directory specified!")
    def __errCb(checkCode, message):
        print(message)
    layout = strict_hdds.get_storage_layout(sys.argv[2])
    layout.check(auto_fix=False, error_callback=__errCb)
    if layout.name in ["bios-ext4", "efi-ext4"]:
        layout.opt_check("swap", auto_fix=False, error_callback=__errCb)
    if layout.name == "efi-bcache-btrfs":
        layout.opt_check("bcache-write-mode", "writeback", auto_fix=False, error_callback=__errCb)
    if True:
        layout.opt_check("mount-write-mode", auto_fix=False, error_callback=__errCb)


def doUmount():
    if len(sys.argv) < 3:
        die("No mount point directory specified!")
    layout = strict_hdds.get_storage_layout(sys.argv[2])
    layout.umount_and_dispose()


def _show(layout):
    if layout is None:
        print("Storage layout: empty")
        return

    if layout.name == "efi-bcachefs":
        if layout.get_ssd() is not None:
            ssdStr = layout.get_ssd()
            bootDiskStr = ""
        else:
            ssdStr = "None"
            bootDiskStr = " (boot disk: %s)" % (layout.boot_disk)
        print("Storage layout: %s, SSD: %s, HDDs: %s%s" % (layout.name, ssdStr, " ".join(layout.get_hdd_list()), bootDiskStr))
        print("Boot directory: %s" % ("read-write" if layout.get_bootdir_rw_controller().is_writable() else "read-only"))
        return

    if layout.name == "efi-bcache-btrfs":
        if layout.get_ssd() is not None:
            ssdStr = layout.get_ssd()
            bootDiskStr = ""
        else:
            ssdStr = "None"
            bootDiskStr = " (boot disk: %s)" % (layout.boot_disk)
        print("Storage layout: %s, SSD: %s, HDDs: %s%s" % (layout.name, ssdStr, " ".join(layout.get_hdd_list()), bootDiskStr))
        print("Boot directory: %s" % ("read-write" if layout.get_bootdir_rw_controller().is_writable() else "read-only"))
        return

    if layout.name == "efi-btrfs":
        if layout.get_ssd() is not None:
            bootDiskStr = ""
        else:
            bootDiskStr = " (boot disk: %s)" % (layout.boot_disk)
        print("Storage layout: %s, HDDs: %s%s" % (layout.name, " ".join(layout.get_hdd_list()), bootDiskStr))
        print("Boot directory: %s" % ("read-write" if layout.get_bootdir_rw_controller().is_writable() else "read-only"))
        return

    if layout.name == "efi-ext4":
        if layout.has_swap_file():
            swapStr = " (with swap)"
        else:
            swapStr = ""
        print("Storage layout: %s, HDD: %s%s" % (layout.name, layout.boot_disk, swapStr))
        print("Boot directory: %s" % ("read-write" if layout.get_bootdir_rw_controller().is_writable() else "read-only"))
        return

    if layout.name == "bios-ext4":
        if layout.has_swap_file():
            swapStr = " (with swap)"
        else:
            swapStr = ""
        print("Storage layout: %s, HDD: %s%s" % (layout.name, layout.boot_disk, swapStr))
        return

    assert False


def _getMountDeviceForPath(pathname):
    buf = subprocess.check_output("mount", text=True)
    for line in buf.split("\n"):
        m = re.search("^(.*) on (.*) type ", line)
        if m is not None and m.group(2) == pathname:
            return m.group(1)
    return None


if __name__ == '__main__':
    if len(sys.argv) <= 1:
        usage()
        sys.exit(1)

    if sys.argv[1] == "--help":
        usage()
        sys.exit(0)

    if os.getuid() != 0:
        die("You must run this command as root!")

    op = sys.argv[1]
    if op == "show":
        doShow()
    elif op == "mount":
        doMount()
    elif op == "create-and-mount":
        doCreateAndMount()
    elif op == "add-harddisk":
        doAddHarddisk()
    elif op == "remove-harddisk":
        doRemoveHarddisk()
    elif op == "create-swap-file":
        doCreateSwapFile()
    elif op == "remove-swap-file":
        doRemoveSwapFile()
    elif op == "create-snapshot":
        doCreateSnapshot()
    elif op == "remove-snapshot":
        doRemoveSnapshot()
    elif op == "sync-from-snapshot":
        doSyncFromSnapshot()
    elif op == "check":
        doCheck()
    elif op == "umount":
        doUmount()
    else:
        die('Invalid operation.')
